import torch 
from typing import List, Dict, Tuple
from pathlib import Path 
import math 
import numpy as np
import scipy.io as scio
import os

def load_block_data(path: str):
    ## 加载原始数据
    data = scio.loadmat(path)
    matrix = data["data"]
    i = 0
    src_data = []
    tgt_data = []
    flag = 0
    while (True):
        if float(matrix[64][i]) == 1.0 or float(matrix[64][i]) == 2.0:
            # 说明想象要开始了
            flag = 1
            save_feature = [matrix[:59, i]]
            tgt_data.append(int(matrix[64, i]))
            i += 1
            continue
        if float(matrix[64, i]) == 251.0:
            # 证明trial结束了～
            save_feature.append(matrix[:59, i])
            src_data.append(save_feature)
            i += 1
            flag = 0
            continue
        if flag == 1:
            ## 依然在想象中
            save_feature.append(matrix[:59, i])
            i += 1
            continue
        if float(matrix[64, i]) == 253.0:
            break
        i += 1

    return src_data, tgt_data

def get_all_data(path):
    ## path = "./data"
    src_data = []
    tgt_data = []
    dirs_files = os.listdir(path)
    for dir_ in dirs_files:
        # print(dir_)
        if os.path.isdir(path + "/" + dir_):
            # print(dir_)
            # 说明是文件夹 那么还要进入此文件夹去遍历所有文件
            files = os.listdir(path + "/" + dir_)
            for file1 in files:
                if file1[-3:] != "mat":
                    continue
                data_file_path = path + "/" + dir_ + "/" + file1
                # print(data_file_path)
                src_block_data, tgt_block_data = load_block_data(data_file_path)
                src_data.extend(src_block_data)
                tgt_data.extend(tgt_block_data)
    
    return src_data, tgt_data

def normalize_data(data):
    #数据归一化
    max_v = torch.max(data)
    min_v = torch.min(data)
    mean_v = torch.mean(data)
    data = 10 * (data - mean_v) / (max_v - min_v)
    data = data - torch.mean(data)
    return data 

def standard_data(data):
    ## 数据标准化 均值为0， 方差为1
    data = (data - torch.mean(data)) / torch.std(data)
    return data
    
if __name__ == "__main__":
    pass
    # data = scio.loadmat("./mi_TrainData/S1/block1.mat")
    # data = data["data"]
    # print(data.shape)
    # for i in range(0,60500):
    #     label = data[64, i]
    #     if label == 1.0 or label == 2.0 or label == 251.0 or label == 253.0 or label == 252.0:
    #         print(i)
    #         print(label)

    #### 保存原始数据
    # src_data, tgt_data = get_all_data("./mi_TrainData")
    # # print(torch.tensor(src_data).shape)
    # print(torch.tensor(src_data).shape)
    # # tgt_data = torch.tensor(tgt_data, dtype=torch.long) - 1
    # src_data = torch.tensor(src_data, dtype=torch.float32)
    # src_data = src_data[:, :375, :]
    # print(src_data.shape)
    # torch.save(src_data, "./train_src_data_raw_375_b.pkl")
    # torch.save(tgt_data, "./train_tgt_data_raw_b.pkl")

     #### 保存测试数据
    # src_data, tgt_data = get_all_data("./mi_TrainData测试数据")
    # # print(torch.tensor(src_data).shape)
    # print(torch.tensor(src_data).shape)
    # # tgt_data = torch.tensor(tgt_data, dtype=torch.long) - 1
    # src_data = torch.tensor(src_data, dtype=torch.float32)
    # src_data = src_data[:, :375, :]
    # print(src_data.shape)
    # torch.save(src_data, "./test_src_data_raw_375_b.pkl")
    # torch.save(tgt_data, "./test_tgt_data_raw_b.pkl")


    # ## 测试特征数据
    # src_data, tgt_data = get_all_feature_data("./data_060520")
    # print(src_data.shape)
    # print(tgt_data.shape)
    # torch.save(src_data, "./psd_feature_data.pkl")
    # torch.save(tgt_data, "./psd_tgt_data.pkl")
    # src = torch.load("./src_data.pkl")
    # print(src[5])


    ## 使用29 30 两个通道的数据 进行训练试试
    # src_data, tgt_data = get_all_data("./data")
    # print(torch.tensor(src_data).shape)
    # print(torch.tensor(tgt_data).shape)
    # src_data = torch.tensor(src_data, dtype=torch.float32)
    # src_data = normalize_data(src_data) # 数据归一化
    # torch.save(src_data, "./src_data_mean_29_30.pkl")
    # torch.save(torch.tensor(tgt_data, dtype=torch.long), "./tgt_data_raw.pkl")

    # s3_train_src_1, s3_train_tgt_1 = load_block_data("./S3/block1.mat")

    # s3_train_src_2, s3_train_tgt_2 = load_block_data("./S3/block2.mat")

    # s3_train_src_1.extend(s3_train_src_2)
    # s3_train_tgt_1.extend(s3_train_tgt_2)

    # s3_train_src = torch.tensor(s3_train_src_1, dtype=torch.float32)[:, :375, :]
    # s3_train_tgt = torch.tensor(s3_train_tgt_1) - 1

    # torch.save(s3_train_src, "./s3_train_src.pkl")
    # torch.save(s3_train_tgt, "./s3_train_tgt.pkl")


    s3_test_src_1, s3_test_tgt_1 = load_block_data("./S3/block3.mat")

    s3_test_src = torch.tensor(s3_test_src_1, dtype=torch.float32)[:, :375, :]
    s3_test_tgt = torch.tensor(s3_test_tgt_1) - 1

    print(s3_test_src.shape)

    torch.save(s3_test_src, "./s3_test_src.pkl")
    torch.save(s3_test_tgt, "./s3_test_tgt.pkl")

    
